"""
系统资源管理模块
包含资源清理、信号处理等系统级功能
"""
import os
import platform
import signal
import atexit
import torch


def cleanup_resources():
    """清理资源函数"""
    try:
        # 清理 PyTorch 缓存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        elif torch.backends.mps.is_available():
            # 检查是否有 mps.empty_cache 方法
            if hasattr(torch, 'mps') and hasattr(torch.mps, 'empty_cache'):
                torch.mps.empty_cache()
        
        print("🧹 资源清理完成")
    except Exception as e:
        print(f"⚠️ 清理资源时出错: {e}")


def signal_handler(signum, frame):
    """信号处理函数"""
    print(f"\n📡 接收到信号 {signum}")
    try:
        cleanup_resources()
    except:
        pass
    os._exit(0)


def setup_macos_multiprocessing():
    """设置macOS多进程环境"""
    if platform.system() == 'Darwin':  # macOS
        import multiprocessing
        # 设置环境变量，这对于避免在fork()后出现Metal相关的崩溃至关重要
        os.environ['OBJC_DISABLE_INITIALIZE_FORK_SAFETY'] = 'YES'
        
        # 强制使用 'spawn' 启动方法，这是在macOS上使用PyTorch和多进程时最安全的方式
        try:
            if multiprocessing.get_start_method() != 'spawn':
                multiprocessing.set_start_method('spawn', force=True)
            print("✅ macOS: Multiprocessing start method set to 'spawn'.")
        except RuntimeError:
            # 如果已经设置过，忽略错误
            print("ℹ️ macOS: Multiprocessing start method was already set.")
            pass


def register_cleanup_handlers():
    """注册清理函数和信号处理"""
    atexit.register(cleanup_resources)
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)


def print_startup_info(device, workers_num):
    """打印启动信息"""
    print(f"🚀 ResNeXt 训练平台启动")
    print(f"💻 运行平台: {platform.system()} {platform.machine()}")
    print(f"🔧 计算设备: {device}")
    print(f"👥 Worker 数量: {workers_num}")
    
    if device.type == 'mps':
        print("🍎 检测到 Apple Silicon，使用 MPS 加速")
    elif device.type == 'cuda':
        print(f"🎮 检测到 CUDA，GPU: {torch.cuda.get_device_name()}")
    else:
        print("💻 使用 CPU 计算")


def get_server_settings(env_name):
    """获取服务器启动设置"""
    is_development = env_name.lower() == 'development'
    
    if is_development:
        print("🔧 开发模式：启用debug模式")
        debug_flag = True
        use_reloader_flag = False if platform.system() == 'Darwin' else True
    else:
        print("🔒 生产模式：禁用debug和reloader")
        debug_flag = False
        use_reloader_flag = False
    
    print(f"🔧 服务器设置: debug={debug_flag}, use_reloader={use_reloader_flag}")
    
    return debug_flag, use_reloader_flag
